import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

# 1. 读取处理后的数据
result = pd.read_excel('monthly_average_result.xlsx')

# 2. 筛选2016-2021年数据并按时间排序
data = result[(result['年'] >= 2016) & (result['年'] <= 2021)]
data = data.sort_values(by=['年', '月'])
time_series = data['平均流量(m3/s)'].values  # 提取流量序列
n = len(time_series)  # 序列长度（72个月）

# 3. 计算UF统计量（正向序列）
UF = np.zeros(n)
for k in range(n):
    x = time_series[:k+1]
    R = np.sum([1 for i in range(1, len(x)) for j in range(i) if x[i] > x[j]])
    E_S = len(x) * (len(x) + 1) / 4
    Var_S = len(x) * (len(x) - 1) * (2 * len(x) + 5) / 72
    UF[k] = (R - E_S) / np.sqrt(Var_S) if Var_S != 0 else 0  # 避免方差为0

# 4. 计算UB统计量（逆向序列）
reverse_series = time_series[::-1]
UB = np.zeros(n)
for k in range(n):
    x = reverse_series[:k+1]
    R = np.sum([1 for i in range(1, len(x)) for j in range(i) if x[i] > x[j]])
    E_S = len(x) * (len(x) + 1) / 4
    Var_S = len(x) * (len(x) - 1) * (2 * len(x) + 5) / 72
    UB[k] = -(R - E_S) / np.sqrt(Var_S) if Var_S != 0 else 0  # 取反并反转顺序
UB = UB[::-1]  # 反转回原时间顺序，使UB与UF对应

# 5. 绘制M-K突变检验曲线
plt.figure(figsize=(10, 5))
plt.plot(UF, label='UFk', color='red')
plt.plot(UB, label='UBk', color='blue')
plt.axhline(y=1.96, linestyle='--', color='gray', alpha=0.7, label='95%置信临界线')
plt.axhline(y=-1.96, linestyle='--', color='gray', alpha=0.7)

# 标注时间轴（每12个月标注年份，共6年）
xticks_pos = np.arange(n)
xticks_labels = [f'{year}-{month:02d}' for year, month in zip(data['年'], data['月'])]
plt.xticks(xticks_pos[::12], labels=xticks_pos[::12]//12 + 2016)  # 每12个月标注年份（2016-2021）
plt.xlabel('时间（年份）')
plt.ylabel('M-K统计量')
plt.title('2016-2021年月均流量突变性检验')
plt.legend()
plt.grid(alpha=0.3)

# 保存图片（与图8格式一致）
plt.savefig('图8. 2016-2021年月均流量突变性检验.png', dpi=300, bbox_inches='tight')
plt.show()